import datetime
import jqdatasdk
import pandas as pd
from stocksys.quant_tactics.UserControl import SimulateControl
from stocksys import models as stocksys_models


class StockInfo:
    current_price = 0
    order_price = 0
    stock_code = ""
    number = 0
    data = None

    def __init__(self, stock_code):
        self.stock_code = stock_code

    def init_data(self, start_date: datetime.datetime, end_date: datetime.datetime, method="JQDATA"):
        if method == "JQDATA":
            security = jqdatasdk.get_security_info(self.stock_code)
            data = jqdatasdk.get_price(security, start_date=start_date, end_date=end_date,
                                       frequency='daily',
                                       fields=None,
                                       skip_paused=False, fq='pre')
            data['date'] = data.index
            rows, cols = data.shape
            data.index = [x for x in range(rows)]
            self.data = data


# 回测模块
class RuleExcuter:
    init_function = None
    before_market_open = None
    market_open = None
    after_market_close = None
    connect_model = None
    context = None
    time_now = None
    data = None
    industry_index = None  # 行业指标
    market_index = None  # 大盘指标
    concept_index = None  # 概念指标

    def __init__(self, quant_tatics_name, context: SimulateControl):
        quant_tatics = __import__(quant_tatics_name)
        init_function = quant_tatics.__getattribute__(quant_tatics_name).init_function
        before_market_open = quant_tatics.__getattribute__(quant_tatics_name).before_market_open
        market_open = quant_tatics.__getattribute__(quant_tatics_name).market_open
        connect_model = quant_tatics.__getattribute__(quant_tatics_name).connect_model
        after_market_close = quant_tatics.__getattribute__(quant_tatics_name).after_market_close

        self.context = context
        self.init_function = init_function
        self.before_market_open = before_market_open
        self.market_open = market_open
        self.connect_model = connect_model
        self.after_market_close = after_market_close

    def get_data(self, security, count, unit, fields):
        pass

    # 获取当前涨跌幅
    def get_growth_rate(self, security):
        last_time = self.time_now + datetime.timedelta(days=-1)
        last_time_rows = self.data[self.data['date'] == last_time]
        current_time_rows = self.data[self.data['date'] == self.time_now]
        current_chg = (current_time_rows['open'] - last_time_rows['close']) / current_time_rows['open'] * 100

    # 获取概念数据
    def get_conceptual(self, start_date, count, security):
        self.start_request()
        data = jqdatasdk.attribute_history_engine(end_dt=start_date, count=count, security=security, unit='1d',
                                                  fields=(
                                                      'open', 'close', 'high', 'low', 'volume', 'money'),
                                                  skip_paused=True,
                                                  df=True, fq=None)
        self.data = data
        self.end_request()

    def MACD(self, short_, long_, m):
        new_data = self.data['date']
        new_data['diff'] = self.data['close'].ewm(adjust=False, alpha=2 / (short_ + 1), ignore_na=True).mean() - \
                           self.data['close'].ewm(adjust=False, alpha=2 / (long_ + 1), ignore_na=True).mean()
        new_data['dea'] = new_data['diff'].ewm(adjust=False, alpha=2 / (m + 1), ignore_na=True).mean()
        new_data['macd'] = 2 * (new_data['diff'] - new_data['dea'])
        return new_data

    def FFT(self):
        pass

    # 获取行业数据
    def get_industry(self):
        self.start_request()

        self.end_request()

    # 获取大盘数据
    def get_marlet_index(self):
        self.start_request()

        self.end_request()

    def init_base_data(self, fullcode, start_date: datetime.datetime, end_date: datetime.datetime, type="stock",
                       method="JQDATA"):
        self.now_date = start_date

    def init_data_from_local(self, path):
        self.data = pd.read_csv(path)

    def default_init_function(self, context):
        return context

    def default_before_market_open(self, context):
        return context

    def default_after_market_close(self, context):
        return context

    def default_market_open(self, context):
        return context

    def default_connect_model(self, context):
        return context

    def default_output_function(self, context):
        return context

    def init_data(self, start_date: datetime.datetime, end_date: datetime.datetime, context: SimulateControl,
                  run_type='daily', method="JQDATA"):
        for ids in range(len(context.select_stock)):
            _stock = context.select_stock[ids]
            _stock = StockInfo(_stock)
            _stock.init_data(start_date, end_date)
            context.select_stock[ids] = _stock
        return context

    def init_excute(self):
        MA30 = self.MACD(30, 20, 10)
        MA90 = self.MACD(30, 20, 10)
        MA5 = self.MACD(30, 20, 10)
        MA3 = self.MACD(30, 20, 10)
        # TODO:调用初始化函数
        return self.init_function(self.context)

    def excute(self, start_date, end_date, run_type="daily"):
        # time_now 全局运行指针，指向当前时间
        context = self.init_excute()

        if run_type == 'daily':
            count_daily = abs(
                (datetime.datetime.strptime(start_date, "%Y-%m-%d") - datetime.datetime.strptime(end_date,
                                                                                                 "%Y-%m-%d")).days)
            context = self.init_data(start_date=start_date, end_date=end_date, context=context, run_type='daily')
            for x in range(count_daily):
                context = self.run_once(context)
            context.get_report()

    # 运行一次
    def run_once(self, context):
        # TODO:调用盘前函数 刚9点时
        context = self.before_market_open(context)
        # TODO:调用盘中函数 9点后
        context = self.market_open(context)
        # TODO:调用盘后函数
        context = self.after_market_close(context)
        # 结算收益
        return context

    def start_request(self):
        api_keys = stocksys_models.APIKeys.objects.all().first()
        jqdatasdk.auth(api_keys.username, api_keys.password)

    def end_request(self):
        jqdatasdk.logout()
